package stream

import "go-tools/run_time"

type Consumer func(interface{})

type FilterFunc func(interface{}) bool
type MapFunc func(interface{}) interface{}
type KeyFunc func(interface{}) string
type MatchFunc func(interface{}) bool

func Range(source <-chan interface{}) Stream {
	p := &ReferencePipeline{
		Prev:        nil,
		SourceStage: nil,
		Next:        nil,
		Source:      source,
		Depth:       0,
	}
	p.SourceStage = p
	return p
}

func Just(items ...interface{}) Stream {
	source := make(chan interface{}, len(items))
	for _, item := range items {
		source <- item
	}
	close(source)
	return Range(source)
}

type Stream interface {
	Filter(FilterFunc) Stream
	Map(MapFunc) Stream
	Distinct(KeyFunc) Stream

	ForEach(Consumer)
	Count() int
	Collect() []interface{}
	AnyMatch(MatchFunc) bool
	AllMatch(MatchFunc) bool
}

type Sink struct {
	MatchRet bool
	Total    int
	Rets     []interface{}

	Name                  string
	DistinctMap           map[string]interface{}
	Downstream            *Sink
	Begin                 func(int64)
	Accept                func(interface{})
	End                   func()
	CancellationRequested func() bool
	GetRet                func() interface{}
}

func NewChainedSink(name string) *Sink {
	s := new(Sink)
	s.Name = name
	s.Begin = func(i int64) {
		if s.Downstream != nil && s.Downstream.Begin != nil {
			s.Downstream.Begin(i)
		}
	}
	s.End = func() {
		if s.Downstream != nil && s.Downstream.End != nil {
			s.Downstream.End()
		}
	}
	s.CancellationRequested = func() bool {
		if s.Downstream != nil && s.Downstream.CancellationRequested != nil {
			return s.Downstream.CancellationRequested()
		}
		return false
	}
	return s
}

type ReferencePipeline struct {
	Prev        *ReferencePipeline
	SourceStage *ReferencePipeline
	Next        *ReferencePipeline
	Source      <-chan interface{}
	Depth       int32
	OpSink      *Sink
}

func (rp *ReferencePipeline) Op(s *Sink) Stream {
	p := &ReferencePipeline{
		Prev:        rp,
		SourceStage: rp.SourceStage,
		Depth:       rp.Depth + 1,
	}
	p.Prev.Next = p
	p.OpSink = s
	return p
}

func (rp *ReferencePipeline) Apply(s *Sink) *Sink {
	var s1 = s
	for p := rp; p.Depth > 0; {
		p.OpSink.Downstream = s
		s = p.OpSink
		p = p.Prev
	}
	size := int64(len(rp.SourceStage.Source))
	if s.Begin != nil {
		s.Begin(size)
	}
	for item := range rp.SourceStage.Source {
		run_time.SafeRun(func() {
			s.Accept(item)
		})
		if s.CancellationRequested() {
			break
		}
	}
	if s.End != nil {
		s.End()
	}
	return s1
}

func (rp *ReferencePipeline) Filter(ff FilterFunc) Stream {
	s := NewChainedSink("filter")
	s.Accept = func(i interface{}) {
		if ff(i) {
			s.Downstream.Accept(i)
		}
	}
	return rp.Op(s)
}

func (rp *ReferencePipeline) Map(mf MapFunc) Stream {
	s := NewChainedSink("map")
	s.Accept = func(item interface{}) {
		s.Downstream.Accept(mf(item))
	}
	return rp.Op(s)
}

func (rp *ReferencePipeline) Distinct(kf KeyFunc) Stream {
	s := NewChainedSink("distinct")
	s.Begin = func(i int64) {
		s.DistinctMap = make(map[string]interface{})
		if s.Downstream != nil && s.Downstream.Begin != nil {
			s.Downstream.Begin(i)
		}
	}
	s.Accept = func(item interface{}) {
		k := kf(item)
		if s.DistinctMap[k] == nil {
			s.DistinctMap[k] = item
			s.Downstream.Accept(item)
		}
	}
	s.End = func() {
		s.DistinctMap = nil
		if s.Downstream != nil && s.Downstream.End != nil {
			s.Downstream.End()
		}
	}
	return rp.Op(s)
}

func (rp *ReferencePipeline) AnyMatch(mf MatchFunc) bool {
	s := new(Sink)
	s.Name = "anyMatch"
	s.Accept = func(item interface{}) {
		if mf(item) {
			s.CancellationRequested = func() bool {
				return true
			}
			s.MatchRet = true
		}
	}
	s.GetRet = func() interface{} {
		return s.MatchRet
	}
	ret := rp.Apply(s)
	return ret.GetRet().(bool)
}

func (rp *ReferencePipeline) AllMatch(mf MatchFunc) bool {
	s := new(Sink)
	s.Name = "allMatch"
	s.Accept = func(item interface{}) {
		if !mf(item) {
			s.CancellationRequested = func() bool {
				return true
			}
			s.MatchRet = false
		}
	}
	s.MatchRet = true
	s.CancellationRequested = func() bool {
		return false
	}
	s.GetRet = func() interface{} {
		return s.MatchRet
	}
	ret := rp.Apply(s)
	return ret.GetRet().(bool)
}

func (rp *ReferencePipeline) ForEach(c Consumer) {
	rp.Apply(&Sink{
		Name: "forEach",
		Accept: func(item interface{}) {
			c(item)
		},
	})
}

func (rp *ReferencePipeline) Count() int {
	s := &Sink{
		Name: "count",
	}
	s.Accept = func(item interface{}) {
		s.Total++
	}
	s.GetRet = func() interface{} {
		return s.Total
	}
	ret := rp.Apply(s)
	return ret.GetRet().(int)
}

func (rp *ReferencePipeline) Collect() []interface{} {
	s := &Sink{
		Name: "collect",
	}
	s.Begin = func(size int64) {
		var r []interface{}
		s.Rets = r
	}
	s.Accept = func(item interface{}) {
		s.Rets = append(s.Rets, item)
	}
	s.GetRet = func() interface{} {
		return s.Rets
	}
	ret := rp.Apply(s)
	return ret.GetRet().([]interface{})
}
